import numpy as np


class SampleTime:
    def __init__(self, result_n_frames: int = 16) -> None:
        self.result_n_frames = result_n_frames

    def __call__(self, image_data: dict) -> np.ndarray:
        image = image_data['image']
        es_time_step = image_data['es_time_step']
        total_frames = image.shape[0]

        es_time_step = min(max(es_time_step, 0), total_frames - 1)

        frames_to_keep = set()
        frames_to_keep.add(0)
        frames_to_keep.add(es_time_step)

        if es_time_step + 1 < total_frames:
            frames_to_keep.add(es_time_step + 1)
            sampling_end = es_time_step + 1
        else:
            sampling_end = es_time_step 

        num_frames_to_sample = self.result_n_frames - len(frames_to_keep)

        if num_frames_to_sample > 0:
            candidate_frames = [i for i in range(1, sampling_end + 1) if i not in frames_to_keep]

            if len(candidate_frames) < num_frames_to_sample:
                additional_needed = num_frames_to_sample - len(candidate_frames)
                extra_frames = range(sampling_end + 1, min(total_frames, sampling_end + 1 + additional_needed))
                candidate_frames.extend(extra_frames)

            if candidate_frames:
                sampled_indices = np.linspace(
                    0, len(candidate_frames) - 1, num=min(num_frames_to_sample, len(candidate_frames)), dtype=int)
                sampled_frames = [candidate_frames[i] for i in sampled_indices]
                frames_to_keep.update(sampled_frames)

        frames_to_keep = [f for f in frames_to_keep if f < total_frames]
        frames_to_keep = sorted(frames_to_keep)

        while len(frames_to_keep) < self.result_n_frames:
            frames_to_keep.append(total_frames - 1)

        sampled_image = image[frames_to_keep]

        return sampled_image
    

class SampleMRI:
    def __init__(self, result_n_frames):
        self.result_n_frames = result_n_frames

    def __call__(self, image_data: dict) -> np.ndarray:
        image = image_data['image']
        num_slices = image.shape[0]

        indices = np.linspace(0, num_slices - 1, self.result_n_frames, dtype=int)

        sampled_image = image[indices]
        return sampled_image
